#!/bin/bash
#############################################################################
#  Copyright (C) 2013-2015 Lawrence Livermore National Security, LLC.
#  Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
#  Written by Albert Chu <chu11@llnl.gov>
#  LLNL-CODE-644248
#
#  This file is part of Magpie, scripts for running Hadoop on
#  traditional HPC systems.  For details, see https://github.com/llnl/magpie.
#
#  Magpie is free software; you can redistribute it and/or modify it
#  under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  Magpie is distributed in the hope that it will be useful, but
#  WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#  General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with Magpie.  If not, see <http://www.gnu.org/licenses/>.
#############################################################################

# This is used by scripts, don't edit this
#
# Export common magpie exports promised by user documentation.

source ${MAGPIE_SCRIPTS_HOME}/magpie/lib/magpie-lib-helper
source ${MAGPIE_SCRIPTS_HOME}/magpie/lib/magpie-lib-log

__Magpie_calc_slurm_magpie_timelimit () {
    # Do not excessively call this
    if [ "${MAGPIE_TIMELIMIT_MINUTES}X" == "X" ]
    then
        local timelimitwalltime=`squeue -j ${SLURM_JOB_ID} -h -o %l`
        Magpie_walltime_to_minutes ${timelimitwalltime}
        export MAGPIE_TIMELIMIT_MINUTES=${magpie_walltimetominutes}
    fi
}

__Magpie_calc_flux_magpie_timelimit () {
    # Do not call this multiple times if it can be avoided
    jobid=$1
    if [ "${MAGPIE_TIMELIMIT_MINUTES}X" == "X" ]
    then
        # duration stored in seconds
        local duration=`flux --parent jobs -no {duration} ${jobid}`
        export MAGPIE_TIMELIMIT_MINUTES=$(echo "${duration} / 60" | bc)
    fi
}

__Magpie_handle_job_name_corner_cases () {
    # Handle special corner cases where user's job name may not
    # play nice, as we use it in file paths and what not
    if [[ $MAGPIE_JOB_NAME =~ "\$" ]]; then
        echo "Job name contains dollar sign, internally replacing with underscore"
        export MAGPIE_JOB_NAME=${MAGPIE_JOB_NAME//\$/_}
    fi
    if [[ $MAGPIE_JOB_NAME =~ [[:space:]] ]]; then
        echo "Job name contains whitespace, internally replacing with underscore"
        export MAGPIE_JOB_NAME=${MAGPIE_JOB_NAME//[[:space:]]/_}
    fi
}

if [ "${MAGPIE_SUBMISSION_TYPE}" == "sbatchsrun" ]
then
    export MAGPIE_CLUSTER_NODERANK="${SLURM_NODEID:=0}"
    export MAGPIE_NODE_COUNT="${SLURM_NNODES}"
    export MAGPIE_NODELIST="${SLURM_JOB_NODELIST}"
    export MAGPIE_JOB_NAME="${SLURM_JOB_NAME}"
    export MAGPIE_JOB_ID="${SLURM_JOB_ID}"
    __Magpie_calc_slurm_magpie_timelimit
    __Magpie_handle_job_name_corner_cases
elif [ "${MAGPIE_SUBMISSION_TYPE}" == "sbatchmpirun" ]
then
    export MAGPIE_CLUSTER_NODERANK="${SLURM_NODEID:=0}"
    export MAGPIE_NODE_COUNT="${SLURM_NNODES}"
    export MAGPIE_NODELIST="${SLURM_JOB_NODELIST}"
    export MAGPIE_JOB_NAME="${SLURM_JOB_NAME}"
    export MAGPIE_JOB_ID="${SLURM_JOB_ID}"
    __Magpie_calc_slurm_magpie_timelimit
    __Magpie_handle_job_name_corner_cases
elif [ "${MAGPIE_SUBMISSION_TYPE}" == "msubslurmsrun" ]
then
    export MAGPIE_CLUSTER_NODERANK="${SLURM_NODEID:=0}"
    export MAGPIE_NODE_COUNT="${SLURM_NNODES}"
    export MAGPIE_NODELIST="${SLURM_JOB_NODELIST}"
    export MAGPIE_JOB_NAME="${MOAB_JOBNAME}"
    export MAGPIE_JOB_ID="${SLURM_JOB_ID}"
    __Magpie_calc_slurm_magpie_timelimit
    __Magpie_handle_job_name_corner_cases
elif [ "${MAGPIE_SUBMISSION_TYPE}" == "msubtorquepdsh" ]
then
    export MAGPIE_CLUSTER_NODERANK="$(($(echo ${MOAB_NODELIST} | sed 's/&/\n/g' | grep -n ${HOSTNAME} | cut -f1 -d:) - 1))"
    export MAGPIE_NODE_COUNT="${PBS_NUM_NODES}"
    export MAGPIE_NODELIST=$(echo ${MOAB_NODELIST} | sed 's/&/,/g')
    export MAGPIE_JOB_NAME="${MOAB_JOBNAME}"
    export MAGPIE_JOB_ID="${PBS_JOBID}"
    export MAGPIE_TIMELIMIT_MINUTES=$(expr ${PBS_WALLTIME} / 60)
    __Magpie_handle_job_name_corner_cases
elif [ "${MAGPIE_SUBMISSION_TYPE}" == "lsfmpirun" ]
then
    # This compress/expand is to properly sort hosts, can't lexicographically do it
    # i.e. lexicographically node150 < node30
    compressedhosts=`${MAGPIE_SCRIPTS_HOME}/bin/magpie-compress-nodes ${LSB_HOSTS}`
    expandedhosts=`${MAGPIE_SCRIPTS_HOME}/bin/magpie-expand-nodes ${compressedhosts}`
    myhostname=`hostname`
    if [ "${MAGPIE_HOSTNAME_CMD_MAP}X" != "X" ]
    then
        myhostname=`${MAGPIE_HOSTNAME_CMD_MAP} $myhostname`
        if [ $? -ne 0 ]
        then
            echo "Error in MAGPIE_HOSTNAME_CMD_MAP = ${MAGPIE_HOSTNAME_CMD_MAP}"
            exit 1
        fi
    fi

    noderank=0
    for host in $expandedhosts
    do
	if [ "${host}" == "${myhostname}" ]
	then
	    break
	fi
	noderank=`expr ${noderank} + 1`
    done
    export MAGPIE_CLUSTER_NODERANK="$noderank"
    export MAGPIE_NODE_COUNT="$(echo ${LSB_HOSTS} | tr " " "\n" | sort -u | wc -w)"
    export MAGPIE_NODELIST=$(echo ${LSB_HOSTS} | sed 's/ /,/g')
    export MAGPIE_JOB_NAME="${LSB_JOBNAME}"
    export MAGPIE_JOB_ID="${LSB_JOBID}"
    while [[ $(bjobs -l ${LSB_JOBID}) == *"Status <P"* ]]
    do
        # Pending state gives time in hours, mins, or seconds and it
        # seems that there is no way of know which it is giving. By
        # waiting for running status we can get it properly.
        sleep 1
    done

    jobinfo=$(bjobs -l ${LSB_JOBID})
    # Run State
    walltime=$(echo $jobinfo | awk '{for(i=1;i<=NF;i++){ if($i=="RUNLIMIT"){printf "%.0f\n", $((i+1))} } }')
    export MAGPIE_TIMELIMIT_MINUTES=${walltime}
    __Magpie_handle_job_name_corner_cases
elif [ "${MAGPIE_SUBMISSION_TYPE}" == "fluxbatchrun" ]
then
    jobid=$(flux getattr jobid)
    export MAGPIE_CLUSTER_NODERANK=$(flux getattr rank)
    export MAGPIE_NODE_COUNT=`flux --parent jobs -no {nnodes} ${jobid}`
    export MAGPIE_NODELIST=`flux --parent jobs -no {nodelist} ${jobid}`
    export MAGPIE_JOB_NAME=`flux --parent jobs -no {name} ${jobid}`
    # just in case fancy f is used in flux job id, convert to lower case f
    export MAGPIE_JOB_ID=`FLUX_F58_FORCE_ASCII=t flux job id --to=f58 ${jobid}`
    __Magpie_calc_flux_magpie_timelimit ${jobid}
    __Magpie_handle_job_name_corner_cases
else
    Magpie_output_internal_error "Invalid MAGPIE_SUBMISSION_TYPE = ${MAGPIE_SUBMISSION_TYPE}"
fi
